# visualize.py

import os
import time
import pandas as pd
from sqlalchemy import create_engine
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib.font_manager as fm
import warnings


warnings.filterwarnings("ignore")

# ----------------------------
# 配置部分
# ----------------------------
DB_CONFIG = {
    'host': '192.168.37.1',
    'port': 9030,
    'user': 'root',
    'password': 'pass',
    'database': 'demo',
    'charset': 'utf8mb4'
}

TABLE_NAME = "production_line_equipment_usage"
REFRESH_INTERVAL = 10  # 单位：秒，每10秒刷新一次数据


# ----------------------------
# 创建数据库连接引擎
# ----------------------------
def create_db_engine():
    try:
        connection_string = f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
        engine = create_engine(connection_string, echo=False)
        print(" 数据库连接成功")
        return engine
    except Exception as e:
        print(f"数据库连接失败: {e}")
        return None


# ----------------------------
# 从 Doris 表中读取数据
# ----------------------------
def fetch_data(engine):
    query = f"SELECT * FROM {TABLE_NAME};"
    try:
        df = pd.read_sql(query, engine)
        print(f" 已加载 {len(df)} 条记录")
        return df
    except Exception as e:
        print(f" 查询失败: {e}")
        return pd.DataFrame()


# ----------------------------
# 设置中文字体支持
# ----------------------------
def setup_chinese_fonts():
    available_fonts = [f.name for f in fm.fontManager.ttflist]
    chinese_fonts = [
        'SimHei', 'Microsoft YaHei', 'Arial Unicode MS',
        'Noto Sans CJK SC', 'WenQuanYi Micro Hei', 'DejaVu Sans'
    ]
    selected_font = 'DejaVu Sans'
    for font in chinese_fonts:
        if font in available_fonts:
            selected_font = font
            break
    plt.rcParams['font.sans-serif'] = [selected_font, 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
    return selected_font


# ----------------------------
# 绘制可视化图表
# ----------------------------
def plot_data(df):
    if df.empty:
        print(" 没有数据可供可视化")
        return

    plt.style.use('ggplot')
    sns.set_palette("viridis")
    selected_font = setup_chinese_fonts()

    use_english = (selected_font == 'DejaVu Sans')

    main_title = '故障记录实时统计分析' if not use_english else 'Real-time Fault Records Analysis'
    titles = [
        '按产线分布的故障数量' if not use_english else 'Fault Count by Production Line',
        '故障等级分布' if not use_english else 'Fault Level Distribution',
        '每日新增故障趋势' if not use_english else 'Daily Fault Trend',
        'Top 10 故障关键词' if not use_english else 'Top 10 Keywords'
    ]
    labels = {
        'production_line': '产线' if not use_english else 'Production Line',
        'count': '数量' if not use_english else 'Count',
        'date': '日期' if not use_english else 'Date',
        'keyword': '关键词' if not use_english else 'Keyword'
    }

    fig, axes = plt.subplots(2, 2, figsize=(16, 10))
    fig.suptitle(main_title, fontsize=16)

    # 图1: 按产线分布
    if 'productionLine' in df.columns:
        df.groupby('productionLine').size().plot(kind='bar', ax=axes[0, 0])
        axes[0, 0].set_title(titles[0])
        axes[0, 0].set_xlabel(labels['production_line'])
        axes[0, 0].set_ylabel(labels['count'])

    # 图2: 故障等级分布
    if 'fault_level' in df.columns:
        df['fault_level'].value_counts().sort_index().plot(kind='bar', ax=axes[0, 1])
        axes[0, 1].set_title(titles[1])
        axes[0, 1].set_xlabel('等级')
        axes[0, 1].set_ylabel(labels['count'])

    # 图3: 时间趋势
    if 'created_at' in df.columns:
        df['date'] = pd.to_datetime(df['created_at']).dt.date
        df.groupby('date').size().plot(kind='line', marker='o', ax=axes[1, 0])
        axes[1, 0].set_title(titles[2])
        axes[1, 0].set_xlabel(labels['date'])
        axes[1, 0].set_ylabel(labels['count'])

    # 图4: 关键词 Top10
    if 'keywords' in df.columns:
        from collections import Counter
        all_keywords = []
        for keywords in df['keywords']:
            if isinstance(keywords, list):
                all_keywords.extend(keywords)
        keyword_counter = Counter(all_keywords)
        top_keywords = pd.Series(keyword_counter).sort_values(ascending=False).head(10)
        top_keywords.plot(kind='barh', ax=axes[1, 1], color='skyblue')
        axes[1, 1].set_title(titles[3])
        axes[1, 1].set_xlabel(labels['count'])
        axes[1, 1].set_ylabel(labels['keyword'])

    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    plt.show()


# ----------------------------
# 主程序入口
# ----------------------------
def main():
    print("[INFO] 开始实时可视化 Doris 数据...")
    engine = create_db_engine()
    if not engine:
        return

    while True:
        df = fetch_data(engine)
        plt.close('all')  # 清除旧图
        plot_data(df)
        print(f" 等待 {REFRESH_INTERVAL} 秒后刷新数据...\n")
        time.sleep(REFRESH_INTERVAL)


if __name__ == "__main__":
    main()